from __future__ import with_statement
import numpy
from util import writeFile
import subprocess as sp
from os import path
import os
import tempfile
import numpy as np
import platform


osName = platform.system().lower()

def openFile(filePath):
    if osName == 'windows':
        os.startfile( filePath ) #@UndefinedVariable
    elif osName == 'darwin':
        sp.Popen( ['open', filePath ] )
    else:
        sp.Popen( ['xdg-open', filePath ] )


class TexDoc:
    
    def __init__(self, docClass='article'):
        self.packageL = []
        self.docOptionL = []
        self.doc = ''
        self.docClass = docClass
        
    def usepackage(self, package, opt=None ):
        if opt is not None: optStr = "[%s]"%opt
        else:               optStr = ''
        self.packageL.append( '\\usepackage%s{%s}'%(optStr, package) )
            
    def addDocOption(self, *optL ):
        self.docOptionL + optL
    
    def add( self, texObj ):
        if hasattr(texObj, '__tex__'):
            self.doc += texObj.__tex__(self)
        else:
            self.doc += str(texObj)
        
    
    def __str__(self):
        # clean duplicate packages
        pkL= []
        pkSet = set()
        for pk in self.packageL:
            if not pk in pkSet:
                pkL.append(pk)
                pkSet.add(pk)
        return """
\\documentclass[%s]{%s}
%s
\\begin{document}
%s
\\end{document}
"""%(','.join(self.docOptionL), self.docClass, '\n'.join(pkL), self.doc )

    def build(self, cmd='latex', filePrefix='result', folder=None):
        if folder is None:
            self.tmpDir = path.join(tempfile.gettempdir(), 'texToolsBuild')
        else :
            self.tmpDir = folder
        
        if not path.exists(self.tmpDir): os.makedirs( self.tmpDir )    
           
        texPath = path.join( self.tmpDir, '%s.tex'%filePrefix  )
        writeFile( str(self), texPath )
        
        sp.call( [cmd, '--output-directory', self.tmpDir, texPath]  )
        
    def pdf(self):
        self.build('pdflatex')
        openFile( path.join( self.tmpDir, 'result.pdf' ) )
#        sp.Popen( ['xdg-open', path.join( self.tmpDir, 'result.pdf' ) ] )
    
    def dvi(self):
        self.build('latex')
        openFile( path.join( self.tmpDir, 'result.dvi' ) )
        #sp.Popen( ['xdg-open', path.join( self.tmpDir, 'result.dvi' ) ] )



def escapeSpecialChar( _str ):
    charList = ['\\','_','$','{']
    for char in charList:
        _str = _str.replace(char, "\\"+char)
    return _str

def buildTable(strMat, alignStr=None, hlineIdx=[0,float('inf')] ):
    
    # default align info is left justified with lines at both ends
    if alignStr is None:
        alignStr = '|%s|'%('l'*strMat.shape[1])
    
    # generate some 'useless' align information for tex formatting
    alignDict = {'c':str.center, 'l':str.ljust, 'r':str.rjust }
    alignList = alignStr.replace('|', '' )
    if len(alignList) != strMat.shape[1]:
        alignList = 'l'*strMat.shape[1]
    
    # find the max string length of the column to specify the column sizes
    sz = [0]*strMat.shape[1]
    for line in strMat:
        for i,cell in enumerate(line):
            sz[i] = numpy.max( (len(cell), sz[i]) )
    
    # apply modulo on the hlineIdx list
    for i,hline in enumerate( hlineIdx ):
        if hline < 0:    
            hlineIdx[i] = hline % strMat.shape[0]
    

    # format the table
    table = '\\begin{tabular}{%s}\n'%(alignStr)
    tmpLine = ['']*strMat.shape[1]
    for j,line in enumerate(strMat):
        if j in hlineIdx:
            table += '\\hline\n'
        
        for i,cell in enumerate(line):
#            cell = escapeSpecialChar( cell )
            tmpLine[i] = alignDict[alignList[i]](cell,sz[i])
        table += ' & '.join(tmpLine) + ' \\\\\n'

        
    if max( hlineIdx ) > strMat.shape[0]:
        table += '\\hline\n'
        
    table = table.replace( 'nan', ' - ' )
    return table + '\\end{tabular}\n'


def figure( texStr, caption ):
    return """
    \\begin{figure}
    \\caption{%s}
    %s
    \\end{figure}
    """(caption, texStr)

def tableEnv( texStr, caption=None, opt=None, captionUp=True ):
    texL = ["\n\\begin{table}[ht]"]
    if captionUp and caption is not None: 
        texL.append( "\\caption{%s}"%(caption))
    if opt is not None:
        texL.append( opt )
    texL.append( texStr )
    if not captionUp and caption is not None:
        texL.append( "\\caption{%s}"%(caption) )
    texL.append("\\end{table}\n")
        
    return '\n'.join(texL)


def center( texStr ):
    return """\\begin{center}
%s
\\end{center}
"""%texStr


def bold( _str ):
    return '\\textbf{%s}'%_str

def mathMode( _str ):
    return '$ %S $'%_str




emptyD = {}
class TableMap:
    
    def __init__(self, tableD, rowL=None, colL=None, rowSortKey=None, colSortKey=None, noneStr = '---'):
        
        self.noneStr = noneStr
        
        if rowL is None or colL is None:
            rowKey, colKey = zip(*tableD.keys())
            if rowL is None: 
                rowL = list(set(rowKey))
                rowL.sort( key = rowSortKey )
            if colL is None: 
                colL = list(set(colKey)) 
                colL.sort( key = colSortKey )
        self.row2idx = self.buildIdx(rowL)
        self.col2idx = self.buildIdx(colL)
        self.formatInfo()
        
        self.mat = np.zeros( (len(rowL), len(colL)) , np.object )
        self.mat[:,:] = None
        for key,val in tableD.iteritems():
            try:             self.mat[ self.row2idx[key[0]], self.col2idx[key[1]] ] = val
            except KeyError: pass;
    
    def formatInfo( self, _format = {}, align={}, title=None, rowCaption=emptyD.get, colCaption=emptyD.get ):
        self.format = _format
        self.align = align
        self.title = title
        self.colCaption = colCaption
        self.rowCaption = rowCaption
        self.texFormatD = {}
        
    
    def delCol(self, key ):
        idx = self.col2idx.pop(key)
        self.mat = np.delete( self.mat, idx, 1 )
        for key, idx_ in self.col2idx.iteritems():
            if idx_ > idx : self.col2idx[key] -= 1
    
    def insertCol(self, colName, colD, beforeKey ):
        idx = self.col2idx[beforeKey]
        for key,idx_ in self.col2idx.iteritems():
            if idx_ >= idx: self.col2idx[key] += 1
        self.col2idx[colName] = idx
        self.mat = np.insert(self.mat,idx, None, 1 )
        for key, val in colD.iteritems():
            rowIdx = self.row2idx[key]
            self.mat[rowIdx, idx] = val
        
    
    def orderCol( self, colL ):
        self.mat = np.array( [ self[:,col] for col in colL ] ).T
        self.col2idx = self.buildIdx(colL) 
        
    
    def transpose(self):
        self.mat = self.mat.T
        tmp = self.col2idx
        self.col2idx= self.row2idx
        self.row2idx = tmp
    
    def iterCols(self):
        for key,idx in self.col2idx.iteritems():
            yield key, self.mat[:,idx]
    
    def iterRows(self):
        for key,idx in self.row2idx.iteritems():
            yield key, self.mat[idx,:]

    def colKeys(self):
        return self._sortedKeys(self.col2idx)
    
    def rowKeys(self):
        return self._sortedKeys(self.row2idx)
    
    def _sortedKeys(self, key2idx ):
        keyL = [ (idx,key) for key, idx  in key2idx.iteritems()]
        keyL.sort()
        return zip(*keyL)[1] 
    
    def __getitem__(self, argL):
        row, col = argL
        if not isinstance(row, slice): row = self.row2idx[row]
        if not isinstance(col, slice): col = self.col2idx[col]
        return self.mat.__getitem__( (row,col) )
        
    def __setitem__(self, argL, value):
        row, col = argL
        if not isinstance(row, slice): row = self.row2idx[row]
        if not isinstance(col, slice): col = self.col2idx[col]
        return self.mat.__setitem__( (row,col), value )
        
        
    def buildIdx(self, valL):
        d = {}
        for i, val in enumerate(valL):
            d[val] = i
        return d
    
        
    def toStrMat(self, colHeader=True, rowHeader=True):
        # build the columns
        
        colL = [None]*(len(self.col2idx))
        for colKey, colIdx in self.col2idx.iteritems():
            fmt = self.format.get(colKey,'%s')
            
            col = []
            for e in self.mat[:,colIdx]:
                if e is None: col.append(self.noneStr)
                else:         
                    try:
                        col.append( fmt%e )
                    except TypeError:
                        col.append( str(e) )
            
            colL[colIdx] = col 
#            colL[colIdx] = [ fmt%e for e in self.mat[:,colIdx] ]
            
            if colHeader:
                colL[colIdx].insert(0, self.colCaption(colKey, str(colKey)) ) 
        
        if rowHeader:
            rowHead = ['']*( 1+len(self.row2idx))
            for rowKey, rowIdx in self.row2idx.iteritems():
                caption = self.rowCaption( rowKey, str(rowKey) )
                rowHead[ 1+ rowIdx]= caption
            colL.insert(0, rowHead)
            
        
        return colL
    
    def __tex__(self, texDoc=None):
        if texDoc is not None:
            texDoc.usepackage("geometry", "margin=0.5cm,landscape")
            texDoc.usepackage('xcolor','table')
#        texDoc.usepackage('color','usenames,dvipsnames')
        strMat = self.toStrMat()

#        strMat = [ [ texTools.escapeSpecialChar(e) for e in col] for col in strMat  ] 
        
        strMat[0] = [ bold(cell) for cell in strMat[0] ] # bold row header
        strMat = zip(*strMat) # transpose (we have to do it anyway)
        strMat[0] = [ bold(cell) for cell in strMat[0] ] # bold col header
        
        strMat = np.array( strMat )
        
        alignL = [ self.align.get( key, 'c' ) for key in self.colKeys() ]
        alignStr = 'r|'+''.join(alignL) 
#        alignStr = 'r|%s'%('c'*len(self.col2idx))
        texTable = buildTable(strMat,alignStr,[1])
        texTable = center(texTable)
        texTable = tableEnv(texTable, self.title,opt="\\rowcolors{2}{blue!5}{}")
            
        return texTable
    
    def __str__(self):
        colL = self.toStrMat() 
        colL.insert(1,['|']*len(colL[0]))
        # set string size and center the text in each cell
        for i,col in enumerate(colL):
            maxSz = max( [ len(cell) for cell in col ] )

            colL[i] = [ cell.center(maxSz) for cell in col ]
            colL[i].insert(1,'-'*maxSz )
            
            
        # zip(*colL) is equivalent to transpose
        return '\n'.join( [' '.join(row) for row in zip(*colL) ] )
        
    def __tab__(self,sep='\t'):
        colL = self.toStrMat() 
        return '\n'.join( [sep.join(row) for row in zip(*colL) ] )



